// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::{
    cell::Cell,
    ffi::{c_int, c_void},
};

use smallvec::smallvec;
use tracing::{debug, error, info};

use cudax::cublaslt::*;
use xgpu_common::{
    api_name::ApiFuncName,
    ipc::message::{Argument, ArgumentFlag, Request},
};

use crate::{fault_guard::virt, hook::cublaslt::CublasLtApi, thread};

#[derive(Clone, Copy, Debug, Default)]
struct ScaleInfo {
    scale_type: cudaDataType_t,
    pointer_mode: cublasLtPointerMode_t,
}

fn get_thread_scale_info() -> ScaleInfo {
    THREAD_SCALE_INFO.with(|cell| cell.get())
}

fn set_thread_scale_info(info: ScaleInfo) {
    THREAD_SCALE_INFO.with(|cell| cell.set(info));
}

fn reset_thread_scale_info() {
    set_thread_scale_info(ScaleInfo::default());
}

fn set_scale_type(scale_type: cudaDataType_t) {
    THREAD_SCALE_INFO.with(|cell| {
        let mut current = cell.get();
        current.scale_type = scale_type;
        cell.set(current);
    });
}

fn set_pointer_mode(pointer_mode: cublasLtPointerMode_t) {
    THREAD_SCALE_INFO.with(|cell| {
        let mut current = cell.get();
        current.pointer_mode = pointer_mode;
        cell.set(current);
    });
}

thread_local! {
    static THREAD_SCALE_INFO: Cell<ScaleInfo> = Cell::new(ScaleInfo::default());
}

fn get_scale_type_size(scale_type: u32) -> usize {
    #[allow(non_upper_case_globals)]
    match scale_type {
        // 1-byte types
        CUDA_R_8I | CUDA_R_8U | CUDA_C_8I | CUDA_C_8U | CUDA_R_8F_E4M3 | CUDA_R_8F_E5M2 => 1,

        // 2-byte types
        CUDA_R_16F | CUDA_R_16BF | CUDA_R_16I | CUDA_R_16U => 2,

        // 4-byte types (real 32-bit + complex 16-bit)
        CUDA_R_32F | CUDA_R_32I | CUDA_R_32U | CUDA_C_16F | CUDA_C_16BF | CUDA_C_16I
        | CUDA_C_16U => 4,

        // 8-byte types (real 64-bit + complex 32-bit)
        CUDA_R_64F | CUDA_R_64I | CUDA_R_64U | CUDA_C_32F | CUDA_C_32I | CUDA_C_32U => 8,

        // 16-byte types (complex 64-bit)
        CUDA_C_64F | CUDA_C_64I | CUDA_C_64U => 16,

        // 4-bit types and any other unknown type → default to 8
        _ => 8,
    }
}
pub struct CublasLtApiImpl;

impl CublasLtApi for CublasLtApiImpl {
    fn cublasLtCreate(&self, light_handle: *mut cublasLtHandle_t) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtCreate");
        debug!(
            "before: lt_handle ptr:{:p},  *lt_handle:{:?}",
            light_handle,
            unsafe { *light_handle }
        );

        let mut light_handle_usize = unsafe { *light_handle };

        let req = Request::with_args(
            ApiFuncName::FuncCublasltcreate as u64,
            smallvec![Argument::from_mut(
                &mut light_handle_usize,
                ArgumentFlag::ARG_OUT | ArgumentFlag::ARG_VIRT,
            )],
        );
        let request_id = req.request_id();
        let res = thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed");
        debug!(
            "after: lt_handle ptr:{:p},  *lt_handle:{:?}",
            light_handle,
            unsafe { *light_handle }
        );

        unsafe {
            *light_handle = light_handle_usize as cublasLtHandle_t;

            virt::handle_insert(
                *light_handle as *mut c_void,
                *light_handle as *mut c_void,
                0,
            )
            .expect("handle_insert failed");
            virt::req_id_vhandle_insert(request_id, *light_handle as *mut c_void)
                .expect("req_id_vhandle_insert failed");
        }
        res
    }

    fn cublasLtDestroy(&self, light_handle: cublasLtHandle_t) -> cublasStatus_t {
        info!("[--Hooked--] api_name: cublasLtDestroy");

        let light_handle_usize = virt::handle_map(light_handle as *mut c_void)
            .expect("handle_map failed") as cublasLtHandle_t
            as usize;

        let req = Request::with_args(
            ApiFuncName::FuncCublasltdestroy as u64,
            smallvec![Argument::from_ref(
                &light_handle_usize,
                ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
            )],
        );
        thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed")
    }

    fn cublasLtMatmul(
        &self,
        light_handle: cublasLtHandle_t,
        compute_desc: cublasLtMatmulDesc_t,
        alpha: *const c_void,
        a: *const c_void,
        a_desc: cublasLtMatrixLayout_t,
        b: *const c_void,
        b_desc: cublasLtMatrixLayout_t,
        beta: *const c_void,
        c: *const c_void,
        c_desc: cublasLtMatrixLayout_t,
        d: *mut c_void,
        d_desc: cublasLtMatrixLayout_t,
        algo: *const cublasLtMatmulAlgo_t,
        workspace: *mut c_void,
        workspace_size_in_bytes: usize,
        stream: cudaStream_t,
    ) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtMatmul");

        let light_handle_usize = virt::handle_map(light_handle as *mut c_void)
            .expect("handle_map failed") as cublasLtHandle_t
            as usize;
        let compute_desc_usize = virt::handle_map(compute_desc as *mut c_void)
            .expect("handle_map failed") as cublasLtMatmulDesc_t
            as usize;
        let a_usize = virt::handle_map(a as *mut c_void).expect("handle_map failed")
            as *const c_void as usize;
        let a_desc_usize = virt::handle_map(a_desc as *mut c_void).expect("handle_map failed")
            as cublasLtMatrixLayout_t as usize;
        let b_usize = virt::handle_map(b as *mut c_void).expect("handle_map failed")
            as *const c_void as usize;
        let b_desc_usize = virt::handle_map(b_desc as *mut c_void).expect("handle_map failed")
            as cublasLtMatrixLayout_t as usize;
        let c_usize = virt::handle_map(c as *mut c_void).expect("handle_map failed")
            as *const c_void as usize;
        let c_desc_usize = virt::handle_map(c_desc as *mut c_void).expect("handle_map failed")
            as cublasLtMatrixLayout_t as usize;
        let d_usize = virt::handle_map(d).expect("handle_map failed") as usize;
        let d_desc_usize = virt::handle_map(d_desc as *mut c_void).expect("handle_map failed")
            as cublasLtMatrixLayout_t as usize;
        let workspace_usize = virt::handle_map(workspace).expect("handle_map failed") as usize;
        let stream_usize = virt::handle_map(stream as *mut c_void).expect("handle_map failed")
            as cudaStream_t as usize;

        let scale_info = get_thread_scale_info();
        let alpha_usize: usize;
        let beta_usize: usize;

        #[allow(non_upper_case_globals)]
        let (alpha_arg, beta_arg) = match scale_info.pointer_mode {
            CUBLASLT_POINTER_MODE_HOST => {
                debug!("both alpha and beta on host");
                debug!("alpha:{:p}, *f32 {:?}", alpha, unsafe {
                    *(alpha as *const f32)
                });
                debug!("beta:{:p}, *f32 {:?}", beta, unsafe {
                    *(beta as *const f32)
                });
                let scale_type_size = get_scale_type_size(scale_info.scale_type);
                let alpha_slice =
                    unsafe { std::slice::from_raw_parts(alpha.cast::<u8>(), scale_type_size) };
                let beta_slice =
                    unsafe { std::slice::from_raw_parts(beta.cast::<u8>(), scale_type_size) };
                (
                    Argument::from_slice(alpha_slice, ArgumentFlag::ARG_IN),
                    Argument::from_slice(beta_slice, ArgumentFlag::ARG_IN),
                )
            }
            CUBLASLT_POINTER_MODE_DEVICE | CUBLASLT_POINTER_MODE_DEVICE_VECTOR => {
                debug!("both alpha and beta on device");
                alpha_usize =
                    virt::handle_map(alpha as *mut c_void).expect("handle_map failed") as usize;
                beta_usize =
                    virt::handle_map(beta as *mut c_void).expect("handle_map failed") as usize;
                (
                    Argument::from_ref(&alpha_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                    Argument::from_ref(&beta_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                )
            }
            CUBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_ZERO => {
                debug!("alpha on device, beta zero");
                alpha_usize =
                    virt::handle_map(alpha as *mut c_void).expect("handle_map failed") as usize;
                (
                    Argument::from_ref(&alpha_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                    Argument::empty(),
                )
            }
            CUBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_HOST => {
                debug!("alpha on device, beta on host");
                alpha_usize =
                    virt::handle_map(alpha as *mut c_void).expect("handle_map failed") as usize;
                let scale_type_size = get_scale_type_size(scale_info.scale_type);
                let beta_slice =
                    unsafe { std::slice::from_raw_parts(beta.cast::<u8>(), scale_type_size) };
                (
                    Argument::from_ref(&alpha_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                    Argument::from_slice(beta_slice, ArgumentFlag::ARG_IN),
                )
            }
            _ => {
                error!("bad pointer_mode:{}", scale_info.pointer_mode);
                panic!("bad pointer_mode:{}", scale_info.pointer_mode);
            }
        };

        let req = Request::with_args(
            ApiFuncName::FuncCublasltmatmul as u64,
            smallvec![
                Argument::from_ref(
                    &light_handle_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
                ),
                Argument::from_ref(
                    &compute_desc_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
                ),
                alpha_arg,
                Argument::from_ref(&a_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&a_desc_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&b_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&b_desc_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                beta_arg,
                Argument::from_ref(&c_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&c_desc_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&d_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&d_desc_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                unsafe { Argument::from_ptr(algo, ArgumentFlag::ARG_IN) },
                Argument::from_ref(
                    &workspace_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
                ),
                Argument::from_ref(&workspace_size_in_bytes, ArgumentFlag::ARG_IN),
                Argument::from_ref(&stream_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
            ],
        );
        thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed")
    }

    fn cublasLtMatrixLayoutCreate(
        &self,
        mat_layout: *mut cublasLtMatrixLayout_t,
        type_x: cudaDataType,
        rows: u64,
        cols: u64,
        ld: i64,
    ) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtMatrixLayoutCreate");
        debug!("rows:{}, cols:{}, ld:{}", rows, cols, ld);

        let mut mat_layout_usize = unsafe { *mat_layout } as usize;
        let req = Request::with_args(
            ApiFuncName::FuncCublasltmatrixlayoutcreate as u64,
            smallvec![
                Argument::from_mut(
                    &mut mat_layout_usize,
                    ArgumentFlag::ARG_OUT | ArgumentFlag::ARG_VIRT,
                ),
                Argument::from_ref(&type_x, ArgumentFlag::ARG_IN),
                Argument::from_ref(&rows, ArgumentFlag::ARG_IN),
                Argument::from_ref(&cols, ArgumentFlag::ARG_IN),
                Argument::from_ref(&ld, ArgumentFlag::ARG_IN),
            ],
        );
        let request_id = req.request_id();
        let res = thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed");

        unsafe {
            *mat_layout = mat_layout_usize as cublasLtMatrixLayout_t;

            virt::handle_insert(*mat_layout as *mut c_void, *mat_layout as *mut c_void, 0)
                .expect("handle_insert failed");
            virt::req_id_vhandle_insert(request_id, *mat_layout as *mut c_void)
                .expect("req_id_vhandle_insert failed");
        }
        res
    }

    fn cublasLtMatrixLayoutDestroy(&self, mat_layout: cublasLtMatrixLayout_t) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtMatrixLayoutDestroy");

        let mat_layout_usize = virt::handle_map(mat_layout as *mut c_void)
            .expect("handle_map failed") as cublasLtMatrixLayout_t
            as usize;

        let req = Request::with_args(
            ApiFuncName::FuncCublasltmatrixlayoutdestroy as u64,
            smallvec![Argument::from_ref(
                &mat_layout_usize,
                ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
            )],
        );
        thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed")
    }

    fn cublasLtMatmulDescCreate(
        &self,
        matmul_desc: *mut cublasLtMatmulDesc_t,
        compute_type: cublasComputeType_t,
        scale_type: cudaDataType_t,
    ) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtMatmulDescCreate");
        debug!("matmul_desc:{:p}, **{:?}", matmul_desc, unsafe {
            *matmul_desc
        });

        set_scale_type(scale_type);
        let mut matmul_desc_usize = unsafe { *matmul_desc } as usize;

        let req = Request::with_args(
            ApiFuncName::FuncCublasltmatmuldesccreate as u64,
            smallvec![
                Argument::from_mut(
                    &mut matmul_desc_usize,
                    ArgumentFlag::ARG_OUT | ArgumentFlag::ARG_VIRT,
                ),
                Argument::from_ref(&compute_type, ArgumentFlag::ARG_IN),
                Argument::from_ref(&scale_type, ArgumentFlag::ARG_IN),
            ],
        );
        let request_id = req.request_id();
        let res = thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed");
        debug!("matmul_desc:{:p}, **{:x}", matmul_desc, matmul_desc_usize);

        unsafe {
            *matmul_desc = matmul_desc_usize as cublasLtMatmulDesc_t;

            virt::handle_insert(*matmul_desc as *mut c_void, *matmul_desc as *mut c_void, 0)
                .expect("handle_insert failed");
            virt::req_id_vhandle_insert(request_id, *matmul_desc as *mut c_void)
                .expect("req_id_vhandle_insert failed");
        }
        res
    }

    fn cublasLtMatmulDescDestroy(&self, matmul_desc: cublasLtMatmulDesc_t) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtMatmulDescDestroy");
        reset_thread_scale_info();

        let matmul_desc_usize = virt::handle_map(matmul_desc as *mut c_void)
            .expect("handle_map failed") as cublasLtMatmulDesc_t
            as usize;

        let req = Request::with_args(
            ApiFuncName::FuncCublasltmatmuldescdestroy as u64,
            smallvec![Argument::from_ref(
                &matmul_desc_usize,
                ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
            )],
        );
        thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed")
    }

    fn cublasLtMatmulDescSetAttribute(
        &self,
        matmul_desc: cublasLtMatmulDesc_t,
        attr: cublasLtMatmulDescAttributes_t,
        buf: *const c_void,
        size_in_bytes: usize,
    ) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtMatmulDescSetAttribute");
        debug!("matmul_desc:{:p}, ", matmul_desc);

        let matmul_desc_usize = virt::handle_map(matmul_desc as *mut c_void)
            .expect("handle_map failed") as cublasLtMatmulDesc_t
            as usize;

        #[allow(non_upper_case_globals)]
        match attr {
            CUBLASLT_MATMUL_DESC_SCALE_TYPE => {
                if size_in_bytes != std::mem::size_of::<cudaDataType_t>() {
                    error!("MATMUL_DESC_SCALE_TYPE: size_in_bytes not match cudaDataType_t");
                    panic!("MATMUL_DESC_SCALE_TYPE: size_in_bytes not match cudaDataType_t");
                }
                let scale_type = unsafe { *(buf as *const cudaDataType_t) };
                debug!("  scale_type: {:?}", scale_type);
                set_scale_type(scale_type);
            }
            CUBLASLT_MATMUL_DESC_POINTER_MODE => {
                if size_in_bytes != std::mem::size_of::<cublasLtPointerMode_t>() {
                    error!(
                        "MATMUL_DESC_POINTER_MODE: size_in_bytes not match cublasLtPointerMode_t"
                    );
                    panic!(
                        "MATMUL_DESC_POINTER_MODE: size_in_bytes not match cublasLtPointerMode_t"
                    );
                }
                let pointer_mode = unsafe { *(buf as *const cublasLtPointerMode_t) };
                debug!("  pointer_mode: {:?}", pointer_mode);
                set_pointer_mode(pointer_mode);
            }
            _ => {
                debug!("Unhandled attribute: {:?}", attr);
            }
        };

        let mut buf_arg: Argument = Argument::empty();
        let mut is_virt_handle = false;
        let mapped_buf: usize;
        if size_in_bytes == std::mem::size_of::<usize>() {
            let ptr = buf.cast::<usize>();
            let buf_usize = unsafe { ptr.read() };
            debug!("buf as usize 0x{:x}", buf_usize);
            let vh = buf_usize as *mut c_void;
            if virt::is_vhandle_valid(vh) {
                mapped_buf = virt::handle_map(vh).expect("handle_map failed") as usize;
                is_virt_handle = true;
                buf_arg =
                    Argument::from_ref(&mapped_buf, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT);
            }
        }
        if !is_virt_handle {
            debug!("size_in_byte:{}", size_in_bytes);
            let slice = unsafe { std::slice::from_raw_parts(buf.cast::<u8>(), size_in_bytes) };
            debug!("buf as slice {:p}, slice:{:?}", buf, slice);
            buf_arg = Argument::from_slice(slice, ArgumentFlag::ARG_IN);
        }

        let req = Request::with_args(
            ApiFuncName::FuncCublasltmatmuldescsetattribute as u64,
            smallvec![
                Argument::from_ref(
                    &matmul_desc_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
                ),
                Argument::from_ref(&attr, ArgumentFlag::ARG_IN),
                buf_arg,
                Argument::from_ref(&size_in_bytes, ArgumentFlag::ARG_IN),
            ],
        );
        thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed")
    }

    fn cublasLtMatmulPreferenceCreate(
        &self,
        pref: *mut cublasLtMatmulPreference_t,
    ) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtMatmulPreferenceCreate");

        let mut pref_usize = unsafe { *pref } as usize;

        let req = Request::with_args(
            ApiFuncName::FuncCublasltmatmulpreferencecreate as u64,
            smallvec![Argument::from_mut(
                &mut pref_usize,
                ArgumentFlag::ARG_OUT | ArgumentFlag::ARG_VIRT,
            )],
        );
        let request_id = req.request_id();
        let res = thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed");

        unsafe {
            *pref = pref_usize as cublasLtMatmulPreference_t;

            virt::handle_insert(*pref as *mut c_void, *pref as *mut c_void, 0)
                .expect("handle_insert failed");
            virt::req_id_vhandle_insert(request_id, *pref as *mut c_void)
                .expect("req_id_vhandle_insert failed");
        }
        res
    }

    fn cublasLtMatmulPreferenceDestroy(&self, pref: cublasLtMatmulPreference_t) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtMatmulPreferenceDestroy");

        let pref_usize = virt::handle_map(pref as *mut c_void).expect("handle_map failed")
            as cublasLtMatmulPreference_t as usize;

        let req = Request::with_args(
            ApiFuncName::FuncCublasltmatmulpreferencedestroy as u64,
            smallvec![Argument::from_ref(
                &pref_usize,
                ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
            )],
        );
        thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed")
    }

    fn cublasLtMatmulPreferenceSetAttribute(
        &self,
        pref: cublasLtMatmulPreference_t,
        attr: cublasLtMatmulPreferenceAttributes_t,
        buf: *const c_void,
        size_in_bytes: usize,
    ) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtMatmulPreferenceSetAttribute");
        debug!("buf:{:p}", buf);
        debug!("size_in_byte:{}", size_in_bytes);

        let pref_usize = virt::handle_map(pref as *mut c_void).expect("handle_map failed")
            as cublasLtMatmulPreference_t as usize;

        let mut slice = unsafe { std::slice::from_raw_parts(buf.cast::<u8>(), size_in_bytes) };
        debug!("buf:{:p}, slice:{:?}", buf, slice);
        let mut slice_flag = ArgumentFlag::ARG_IN;
        if size_in_bytes == std::mem::size_of::<usize>() {
            // device pointer!
            let mut _slice = [0u8; std::mem::size_of::<usize>()];
            _slice.copy_from_slice(&slice[..std::mem::size_of::<usize>()]);
            debug!(
                "buf->usize, be:{:#x}, le:{:#x}(Little-endian!)",
                usize::from_be_bytes(_slice),
                usize::from_le_bytes(_slice)
            );
            // convert virt_map to &[u8] again!!
            let vh = usize::from_le_bytes(_slice) as *mut c_void;
            if virt::is_vhandle_valid(vh) {
                let mapped_buf = virt::handle_map(usize::from_le_bytes(_slice) as *mut c_void);
                slice = unsafe {
                    std::slice::from_raw_parts(std::ptr::from_ref(&mapped_buf).cast::<u8>(), 8)
                };
                slice_flag = ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT;
            }
        }
        let req = Request::with_args(
            ApiFuncName::FuncCublasltmatmulpreferencesetattribute as u64,
            smallvec![
                Argument::from_ref(&pref_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&attr, ArgumentFlag::ARG_IN),
                Argument::from_slice(slice, slice_flag),
                Argument::from_ref(&size_in_bytes, ArgumentFlag::ARG_IN),
            ],
        );
        thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed")
    }

    fn cublasLtMatmulAlgoGetHeuristic(
        &self,
        light_handle: cublasLtHandle_t,
        operation_desc: cublasLtMatmulDesc_t,
        a_desc: cublasLtMatrixLayout_t,
        b_desc: cublasLtMatrixLayout_t,
        c_desc: cublasLtMatrixLayout_t,
        d_desc: cublasLtMatrixLayout_t,
        preference: cublasLtMatmulPreference_t,
        requested_algo_count: c_int,
        heuristic_results_array: *mut cublasLtMatmulHeuristicResult_t,
        return_algo_count: *mut c_int,
    ) -> cublasStatus_t {
        info!("[Hooked] api_name: cublasLtMatmulAlgoGetHeuristic");

        let light_handle_usize = virt::handle_map(light_handle as *mut c_void)
            .expect("handle_map failed") as cublasLtHandle_t
            as usize;
        let operation_desc_usize = virt::handle_map(operation_desc as *mut c_void)
            .expect("handle_map failed") as cublasLtMatmulDesc_t
            as usize;
        let a_desc_usize = virt::handle_map(a_desc as *mut c_void).expect("handle_map failed")
            as cublasLtMatrixLayout_t as usize;
        let b_desc_usize = virt::handle_map(b_desc as *mut c_void).expect("handle_map failed")
            as cublasLtMatrixLayout_t as usize;
        let c_desc_usize = virt::handle_map(c_desc as *mut c_void).expect("handle_map failed")
            as cublasLtMatrixLayout_t as usize;
        let d_desc_usize = virt::handle_map(d_desc as *mut c_void).expect("handle_map failed")
            as cublasLtMatrixLayout_t as usize;
        let preference_usize = virt::handle_map(preference as *mut c_void)
            .expect("handle_map failed")
            as cublasLtMatmulPreference_t as usize;

        let req = Request::with_args(
            ApiFuncName::FuncCublasltmatmulalgogetheuristic as u64,
            smallvec![
                Argument::from_ref(
                    &light_handle_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
                ),
                Argument::from_ref(
                    &operation_desc_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
                ),
                Argument::from_ref(&a_desc_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&b_desc_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&c_desc_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(&d_desc_usize, ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT),
                Argument::from_ref(
                    &preference_usize,
                    ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
                ),
                Argument::from_ref(&requested_algo_count, ArgumentFlag::ARG_IN),
                unsafe { Argument::from_mut_ptr(heuristic_results_array, ArgumentFlag::ARG_OUT) },
                unsafe { Argument::from_mut_ptr(return_algo_count, ArgumentFlag::ARG_OUT) },
            ],
        );
        thread::context()
            .agent()
            .invoke_api::<cublasStatus_t>(req)
            .expect("call invoke_api failed")
    }
}
